import os
import cv2
import aidlite_gpu
import numpy as np
from utils import preprocess_img, convert_shape

class feature_extracor_network(object):
    def __init__(self, model_path:str, input_shape:tuple, output_shape:tuple, device:int, int8:bool, 
                        input_scale:float or list=None, input_zero_point:float or list=None,
                        output_sacle:float or list=None, output_zero_point:float or list=None,
                        means:list=[0.485, 0.456, 0.406], stds:list=[0.229, 0.224, 0.225]) -> None:
        assert os.path.exists(model_path), f"{model_path} 不存在"
        self.means = means
        self.stds = stds

        self.model = self.bulid_model(model_path, input_shape, output_shape, device, int8)
        self.input_shape = input_shape
        self.output_shape = output_shape
        self.int8 =int8
        self.input_scale = input_scale
        self.input_zero_point = input_zero_point
        self.output_sacle = output_sacle
        self.output_zero_point = output_zero_point

    def bulid_model(self, model_path:str, input_shape:tuple, output_shape:tuple, device=-1, int8=False):
        '''
        构建在aidlux上运行的模型
        :param model_path: 模型路径
        :param input_shape: 模型输入shape, eg: (1, 320, 320, 3)
        :param output_shape: 模型输出shape, eg: (1, 8300, 85), [(1, 520, 520, 85), (1, 32, 32, 85)]
        :param device: 运行设备, -1:cpu, 0:gpu, 1:cpu+gpu, 2:dsp
        :param int8: 是否是int8模型
        :return aidlte_model
        '''
        aidlite = aidlite_gpu.aidlite(1)
        aidlite.FAST_ANNModel(model_path, 
                                convert_shape(input_shape, int8), 
                                convert_shape(output_shape, int8), 
                                4, 
                                device)
        return aidlite
        
    def __call__(self, imgs:list)->np.ndarray:
        outshape = list(self.output_shape)
        outshape[0] = len(imgs)
        features = np.zeros(shape=outshape, dtype=np.float32)
        if self.int8:
            for i in range(len(imgs)):
                input_img = cv2.resize(imgs[i], (self.input_shape[1], self.input_shape[2])).astype(np.float32)
                if self.input_scale is not None or self.input_zero_point is not None:
                    input_img = input_img/255
                if self.input_scale is not None:
                    input_img = input_img/self.input_scale
                if self.input_zero_point is not None:
                    input_img = input_img + self.input_zero_point
                input_img = input_img.astype(np.uint8)
                self.model.setInput_Int8(input_img)
                self.model.invoke()
                feature = self.model.getOutput_Int8().reshape(*self.output_shape)[0].astype(np.float32)

                if self.output_zero_point is not None:
                    feature = feature - self.output_zero_point
                if self.output_sacle is not None:
                    feature = feature * self.output_sacle
                features[i] = feature
        else:
            for i in range(len(imgs)):
                _input = preprocess_img(imgs[i], keep_ratio=False, 
                                                target_shape=(self.input_shape[1], self.input_shape[2]),
                                                div_num=255,
                                                means=self.means,
                                                stds=self.stds)
                self.model.setInput_Float32(_input)
                self.model.invoke()
                features[i] = self.model.getOutput_Float32().reshape(*self.output_shape)
        return features

if __name__ == "__main__":
    model = feature_extracor_network("./models/feature_extractor_osnet.tflite")
    imgs = []
    for i in range(10):
        imgs.append(np.random.randn(128, 128, 3))
    out = model(imgs)
    print(out.shape)